import os

import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Test re-load model again
from mnist_under_pytorch_train import Net


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    all_preds = []
    all_targets = []
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()
            all_preds.extend(pred.cpu().numpy())
            all_targets.extend(target.cpu().numpy())

    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset), accuracy))

    return accuracy, all_preds, all_targets


def visualize_predictions(model, device, test_loader, num_images=10):
    """可视化预测结果，显示原始图像和识别的数字"""
    model.eval()
    fig, axes = plt.subplots(2, 5, figsize=(15, 8))
    axes = axes.ravel()

    # 获取一批测试数据
    dataiter = iter(test_loader)
    images, labels = next(dataiter)

    # 将数据移到设备上并进行预测
    images, labels = images.to(device), labels.to(device)
    outputs = model(images)
    preds = outputs.argmax(dim=1)
    probs = torch.exp(outputs)  # 转换为概率

    # 将图像移回CPU并转换为numpy以进行可视化
    images = images.cpu()

    for i in range(num_images):
        # 显示图像
        img = images[i].squeeze()
        axes[i].imshow(img, cmap='gray')

        # 获取预测概率
        prob = probs[i].max().item() * 100

        # 设置标题显示真实标签和预测标签
        color = 'green' if preds[i].item() == labels[i].item() else 'red'
        axes[i].set_title(f'True: {labels[i].item()}, Pred: {preds[i].item()}\nConfidence: {prob:.1f}%', 
                         color=color, fontsize=12)
        axes[i].axis('off')

    plt.suptitle('MNIST Digit Recognition Results', fontsize=16)
    plt.tight_layout()
    plt.show()


def main():
    # Check if CUDA is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    # Create model instance
    model = Net()

    # Validate model exist
    if not os.path.exists("output/mnist_cnn.pt"):
        print("Model not found, please train first.")
        return

    # Load the model state dict with proper device mapping
    model_state_dict = torch.load("output/mnist_cnn.pt", map_location=torch.device(device))
    model.load_state_dict(model_state_dict)

    # Move model to the appropriate device
    model = model.to(device)

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    mnist_test_dataset = datasets.MNIST(root="dataset/mnist", train=False, download=True, transform=transform)
    test_loader = DataLoader(dataset=mnist_test_dataset, batch_size=1000, shuffle=True)

    # 运行测试并获取结果
    accuracy, preds, targets = test(model, device, test_loader)

    # 可视化一些预测结果
    visualize_predictions(model, device, test_loader)

    print(f"Model accuracy: {accuracy:.2f}%")


if __name__ == '__main__':
    main()